import numpy as np
import cv2


class PoseStraightener:
    def __init__(self, neck_idx=18, left_ankle_idx=15, right_ankle_idx=16, hip_idx=19):
        self.neck_idx = neck_idx
        self.left_ankle_idx = left_ankle_idx
        self.right_ankle_idx = right_ankle_idx
        self.hip_idx = hip_idx
        self.R = None
        self.origin = None  # 旋转中心：臀部


        # 可以根据你的模型自行调整这些索引
        self.lower_body_indices = [13, 14, 15, 16, 20, 21,22,23,24,25]  # 下半身关键点索引
    def init_from_first_frame(self, keypoints3d, extra_deg=0, manual_angle_deg=10):
        neck = np.array(keypoints3d[self.neck_idx])
        left_ankle = np.array(keypoints3d[self.left_ankle_idx])
        right_ankle = np.array(keypoints3d[self.right_ankle_idx])
        hip = np.array(keypoints3d[self.hip_idx])
        self.origin = hip.copy()

        ankle_center = (left_ankle + right_ankle) / 2.0
        body_axis = neck - ankle_center
        body_axis_unit = body_axis / (np.linalg.norm(body_axis) + 1e-8)
        z_axis = np.array([0, 0, 1])
        axis = np.cross(body_axis_unit, z_axis)
        axis = axis / (np.linalg.norm(axis) + 1e-8)

        if manual_angle_deg is not None:
            # 用户指定角度（默认逆时针），绕 axis 旋转
            angle_rad = -np.radians(manual_angle_deg)
        else:
            # 自动计算旋转角度 + 额外角度
            angle_rad = -np.arccos(np.clip(np.dot(body_axis_unit, z_axis), -1.0, 1.0))
            angle_rad -= np.radians(extra_deg)

        # 构造旋转向量并转换为矩阵
        rvec = axis * -angle_rad
        self.R, _ = cv2.Rodrigues(rvec)

        print(f"[INFO] 使用旋转角度：{np.degrees(angle_rad):.2f}° （手动: {manual_angle_deg is not None}）")

    def straighten(self, keypoints3d):
        if self.R is None:
            self.init_from_first_frame(keypoints3d)
        if self.R is None or self.origin is None:
            raise RuntimeError("未初始化旋转。请先用第一帧调用 init_from_first_frame。")

        keypoints3d = np.array(keypoints3d)
        centered = keypoints3d - self.origin
        rotated = (self.R @ centered.T).T
        return rotated + self.origin
    
    def straighten_lower_body(self, keypoints3d):
        if self.R is None:
            self.init_from_first_frame(keypoints3d)
        if self.R is None or self.origin is None:
            raise RuntimeError("未初始化旋转。请先用第一帧调用 init_from_first_frame。")

        keypoints3d = np.array(keypoints3d)
        rotated_keypoints = keypoints3d.copy()

        # 仅对下半身关键点进行旋转
        for idx in self.lower_body_indices:
            centered = keypoints3d[idx] - self.origin
            rotated = self.R @ centered
            rotated_keypoints[idx] = rotated + self.origin

        return rotated_keypoints